#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
完整房间占用预测系统
"天气预报式"房间占用预测 - 连接传感器预测和房间占用预测的完整流程
"""

import pandas as pd
import numpy as np
import os
import sys
import subprocess
import pickle
import warnings
import json
import argparse
from datetime import datetime, timedelta
import tensorflow as tf

warnings.filterwarnings('ignore')
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

class CompleteRoomPredictionSystem:
    """
    完整房间占用预测系统
    第一阶段：预测未来传感器数据
    第二阶段：基于预测的传感器数据预测房间占用人数
    """

    def __init__(self, api_mode=False):
        """初始化系统"""
        self.api_mode = api_mode
        self.sensor_model = None
        self.sensor_scalers = None
        self.sensor_columns = [
            'S1_Temp', 'S2_Temp', 'S3_Temp', 'S4_Temp',
            'S1_Light', 'S2_Light', 'S3_Light', 'S4_Light',
            'S1_Sound', 'S2_Sound', 'S3_Sound', 'S4_Sound',
            'S5_CO2', 'S5_CO2_Slope', 'S6_PIR', 'S7_PIR'
        ]

        self.occupancy_models = {
            'rf': 'predict/predict_rf.py',
            'lstm': 'predict/predict_lstm.py',
            'xgb': 'predict/predict_xgb.py'
        }

    def output_progress(self, stage, progress, message, details=None):
        """输出进度信息"""
        if self.api_mode:
            progress_data = {
                "type": "progress",
                "stage": stage,
                "progress": progress,
                "message": message,
                "timestamp": datetime.now().isoformat()
            }
            if details:
                progress_data["details"] = details
            
            # 输出JSON格式的进度信息
            print("PROGRESS:" + json.dumps(progress_data))
            sys.stdout.flush()
        else:
            print(f"[{progress}%] {message}")

    def output_result(self, success, data=None, error=None):
        """输出最终结果"""
        if self.api_mode:
            result_data = {
                "type": "result",
                "success": success,
                "timestamp": datetime.now().isoformat()
            }
            if success and data:
                result_data["data"] = data
            if not success and error:
                result_data["error"] = error
            
            print("RESULT:" + json.dumps(result_data))
            sys.stdout.flush()
        else:
            if success:
                print("预测完成！")
            else:
                print(f"预测失败: {error}")

    def load_sensor_model(self):
        """加载传感器预测模型"""
        self.output_progress("init", 5, "正在加载传感器预测模型...")
        
        model_path = 'model/saved_models/sensor_predictor.h5'
        scalers_path = 'model/saved_models/sensor_scalers.pkl'

        if not os.path.exists(model_path) or not os.path.exists(scalers_path):
            raise FileNotFoundError("传感器预测模型文件不存在！请先训练模型。")

        try:
            self.sensor_model = tf.keras.models.load_model(model_path)
            self.output_progress("init", 10, "传感器模型加载完成")

            with open(scalers_path, 'rb') as f:
                self.sensor_scalers = pickle.load(f)

            self.output_progress("init", 15, "传感器预测模型准备就绪")

        except Exception as e:
            error_msg = f"传感器模型加载失败: {e}"
            if not self.api_mode:
                print(error_msg)
            raise Exception(error_msg)

    def prepare_time_features(self, datetime_series):
        """准备时间特征"""
        time_features = pd.DataFrame()
        time_features['Hour_sin'] = np.sin(2 * np.pi * datetime_series.dt.hour / 24)
        time_features['Hour_cos'] = np.cos(2 * np.pi * datetime_series.dt.hour / 24)
        time_features['DayOfWeek_sin'] = np.sin(2 * np.pi * datetime_series.dt.dayofweek / 7)
        time_features['DayOfWeek_cos'] = np.cos(2 * np.pi * datetime_series.dt.dayofweek / 7)
        time_features['IsWeekend'] = (datetime_series.dt.dayofweek >= 5).astype(int)

        return time_features.values

    def predict_future_sensors(self, historical_data, num_predictions=96):
        """
        第一阶段：预测未来传感器数据

        参数:
        historical_data: 历史数据DataFrame (至少300条记录)
        num_predictions: 预测的时间点数量

        返回:
        预测的传感器数据DataFrame
        """
        self.output_progress("predict", 20, "开始第一阶段：预测未来传感器数据")

        # 数据预处理
        df = historical_data.copy()
        if 'DateTime' not in df.columns:
            df['DateTime'] = pd.to_datetime(df['Date'] + ' ' + df['Time'], format='%d-%m-%Y %H:%M:%S')

        df = df.sort_values('DateTime').tail(300).reset_index(drop=True)

        if len(df) < 300:
            raise ValueError(f"历史数据不足！需要至少300条记录，但只有{len(df)}条")

        self.output_progress("predict", 25, f"准备历史数据: {len(df)}条记录")

        # 准备输入数据
        sensor_data = df[self.sensor_columns].values.astype(np.float32)
        time_features = self.prepare_time_features(df['DateTime'])
        combined_features = np.hstack([sensor_data, time_features])

        self.output_progress("predict", 30, "正在标准化数据...")
        
        # 标准化
        scaler = self.sensor_scalers['combined_scaler']
        normalized_data = scaler.transform(combined_features)
        input_sequence = normalized_data.reshape(1, 300, -1)

        # 预测
        self.output_progress("predict", 35, f"开始预测未来{num_predictions}个时间点的传感器数据...")
        predictions = self.sensor_model.predict(input_sequence, verbose=0)

        # 反标准化
        denormalized_predictions = scaler.inverse_transform(predictions[0])
        sensor_predictions = denormalized_predictions[:, :len(self.sensor_columns)]

        # 生成未来时间
        last_datetime = df['DateTime'].max()
        future_times = []
        current_time = last_datetime + timedelta(minutes=30)

        for i in range(num_predictions):
            future_times.append(current_time)
            current_time += timedelta(minutes=30)

        # 创建结果DataFrame
        results = []
        for i, future_time in enumerate(future_times):
            row = {
                'Date': future_time.strftime('%d-%m-%Y'),
                'Time': future_time.strftime('%H:%M:%S')
            }

            # 添加传感器预测值 (保留1位小数)
            for j, col in enumerate(self.sensor_columns):
                row[col] = round(float(sensor_predictions[i, j]), 1)

            results.append(row)

        future_sensors_df = pd.DataFrame(results)

        # 应用传感器约束
        future_sensors_df = self.apply_sensor_constraints(future_sensors_df)

        self.output_progress("predict", 45, f"传感器预测完成，预测了{len(future_sensors_df)}个时间点")

        return future_sensors_df

    def apply_sensor_constraints(self, df):
        """应用传感器数据约束"""
        # 温度传感器：15-35°C
        temp_cols = [col for col in df.columns if 'Temp' in col]
        for col in temp_cols:
            df[col] = np.clip(df[col], 15, 35)

        # 光线传感器：0-1000
        light_cols = [col for col in df.columns if 'Light' in col]
        for col in light_cols:
            df[col] = np.clip(df[col], 0, 1000)

        # 声音传感器：0-1
        sound_cols = [col for col in df.columns if 'Sound' in col]
        for col in sound_cols:
            df[col] = np.clip(df[col], 0, 1)

        # CO2传感器：300-1000 ppm
        if 'S5_CO2' in df.columns:
            df['S5_CO2'] = np.clip(df['S5_CO2'], 300, 1000)

        # CO2斜率：-50到50
        if 'S5_CO2_Slope' in df.columns:
            df['S5_CO2_Slope'] = np.clip(df['S5_CO2_Slope'], -50, 50)

        # PIR传感器：0或1
        pir_cols = [col for col in df.columns if 'PIR' in col]
        for col in pir_cols:
            df[col] = np.round(np.clip(df[col], 0, 1)).astype(int)

        return df

    def call_occupancy_model(self, sensor_row, model_type='rf'):
        """
        调用单个房间占用预测模型

        参数:
        sensor_row: 单行传感器数据 (Series)
        model_type: 模型类型 ('rf', 'lstm', 'xgb')

        返回:
        预测的占用人数 (int)
        """
        if model_type not in self.occupancy_models:
            raise ValueError(f"不支持的模型类型: {model_type}")

        script_path = self.occupancy_models[model_type]

        if not os.path.exists(script_path):
            raise FileNotFoundError(f"模型脚本不存在: {script_path}")

        try:
            # 构建输入数据字符串
            input_data = f"{sensor_row['Date']},{sensor_row['Time']}"
            for col in self.sensor_columns:
                input_data += f",{sensor_row[col]}"

            # 调用模型脚本 (修复编码问题)
            result = subprocess.run(
                [sys.executable, script_path],
                input=input_data,
                text=True,
                capture_output=True,
                timeout=30,
                encoding='utf-8',
                errors='ignore'
            )

            if result.returncode == 0:
                # 解析输出，提取预测人数
                output_lines = result.stdout.strip().split('\n')
                for line in output_lines:
                    if '四舍五入后的人数:' in line:
                        numbers = [int(x) for x in line.split() if x.isdigit()]
                        if numbers:
                            return numbers[0]

                # 如果没找到，返回0
                return 0
            else:
                print(f"{model_type}模型调用失败: {result.stderr}")
                return None

        except subprocess.TimeoutExpired:
            print(f"{model_type}模型调用超时")
            return None
        except Exception as e:
            print(f"{model_type}模型调用异常: {e}")
            return None

    def predict_room_occupancy(self, future_sensors_df, models=['rf', 'lstm', 'xgb']):
        """
        第二阶段：基于预测的传感器数据预测房间占用

        参数:
        future_sensors_df: 预测的传感器数据DataFrame
        models: 使用的模型列表

        返回:
        完整预测结果DataFrame
        """
        self.output_progress("predict", 50, "开始第二阶段：预测房间占用人数")
        self.output_progress("predict", 55, f"使用模型: {', '.join(models)}")
        self.output_progress("predict", 60, f"预测{len(future_sensors_df)}个时间点的房间占用")

        results = []

        total_rows = len(future_sensors_df)
        
        for idx, row in future_sensors_df.iterrows():
            # 计算当前进度 (60% - 85%)
            current_progress = 60 + int((idx / total_rows) * 25)
            
            if (idx + 1) % 10 == 0 or idx == 0:
                self.output_progress("predict", current_progress, f"正在预测: {idx+1}/{total_rows} - {row['Date']} {row['Time']}")

            # 基础数据
            prediction_result = row.to_dict()

            # 使用多个模型预测
            occupancy_predictions = {}
            for model_type in models:
                try:
                    occupancy = self.call_occupancy_model(row, model_type)
                    occupancy_predictions[f'{model_type}_prediction'] = occupancy
                except Exception as e:
                    if not self.api_mode:
                        print(f"      {model_type}模型失败: {e}")
                    occupancy_predictions[f'{model_type}_prediction'] = None

            # 计算集成预测 (平均值，忽略None值)
            valid_predictions = [v for v in occupancy_predictions.values() if v is not None]
            if valid_predictions:
                ensemble_prediction = round(np.mean(valid_predictions))
                occupancy_predictions['ensemble_prediction'] = ensemble_prediction
            else:
                occupancy_predictions['ensemble_prediction'] = 0

            # 合并所有结果
            prediction_result.update(occupancy_predictions)
            results.append(prediction_result)

        results_df = pd.DataFrame(results)

        self.output_progress("save", 90, "房间占用预测完成，正在保存结果...")

        return results_df

    def save_results(self, results_df, save_sensors=True, save_occupancy=True):
        """保存预测结果"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # 确保保存目录存在
        save_dir = os.path.join('data', 'predict_results')
        os.makedirs(save_dir, exist_ok=True)

        saved_files = []

        if save_sensors:
            # 保存传感器预测数据
            sensor_cols = ['Date', 'Time'] + self.sensor_columns
            sensors_df = results_df[sensor_cols].copy()

            sensors_file = os.path.join(save_dir, f"predicted_sensors_{timestamp}.csv")
            sensors_df.to_csv(sensors_file, index=False)
            saved_files.append(sensors_file)
            if not self.api_mode:
                print(f"传感器预测数据已保存: {sensors_file}")

        if save_occupancy:
            # 保存完整预测结果
            complete_file = os.path.join(save_dir, f"complete_room_prediction_{timestamp}.csv")
            results_df.to_csv(complete_file, index=False)
            saved_files.append(complete_file)
            if not self.api_mode:
                print(f"完整预测结果已保存: {complete_file}")
            else:
                self.output_progress("save", 95, f"结果已保存: {complete_file}")

        return saved_files

    def print_prediction_summary(self, results_df):
        """打印预测结果统计"""
        print("\n" + "=" * 60)
        print("完整房间占用预测结果统计")
        print("=" * 60)

        # 时间范围
        print(f"预测时间范围: {results_df['Date'].iloc[0]} {results_df['Time'].iloc[0]} 到 {results_df['Date'].iloc[-1]} {results_df['Time'].iloc[-1]}")
        print(f"预测时长: {len(results_df)*0.5:.1f}小时 ({len(results_df)}个时间点)")

        # 传感器统计
        print(f"\n温度预测范围:")
        temp_cols = [col for col in results_df.columns if 'Temp' in col]
        for col in temp_cols:
            print(f"   {col}: {results_df[col].min():.1f}°C - {results_df[col].max():.1f}°C")

        if 'S5_CO2' in results_df.columns:
            print(f"\nCO2预测: {results_df['S5_CO2'].min():.0f} - {results_df['S5_CO2'].max():.0f} ppm")

        # 房间占用统计
        print(f"\n房间占用预测:")

        # 各模型统计
        model_cols = [col for col in results_df.columns if '_prediction' in col]
        for col in model_cols:
            predictions = results_df[col].dropna()
            if len(predictions) > 0:
                model_name = col.replace('_prediction', '').upper()
                print(f"   {model_name}: {predictions.min()}-{predictions.max()}人 (平均: {predictions.mean():.1f}人)")

                # 人数分布
                distribution = predictions.value_counts().sort_index()
                dist_str = ", ".join([f"{int(k)}人({v}次)" for k, v in distribution.items()])
                print(f"      分布: {dist_str}")

        # 关键统计
        if 'ensemble_prediction' in results_df.columns:
            ensemble_preds = results_df['ensemble_prediction']
            print(f"\n集成预测关键统计:")
            print(f"   总预测时长: {len(results_df)*0.5:.1f}小时")
            print(f"   总人时: {ensemble_preds.sum()*0.5:.1f}人·小时")
            print(f"   平均占用: {ensemble_preds.mean():.1f}人")
            print(f"   峰值占用: {ensemble_preds.max()}人")
            print(f"   空置时间: {sum(ensemble_preds == 0)*0.5:.1f}小时 ({sum(ensemble_preds == 0)/len(ensemble_preds)*100:.1f}%)")

        print("=" * 60)

    def complete_prediction(self, csv_file, num_predictions=96, models=['rf', 'lstm', 'xgb'], save_results=True):
        """
        完整预测流程的主函数

        参数:
        csv_file: 历史数据CSV文件路径
        num_predictions: 预测的时间点数量
        models: 使用的房间占用预测模型
        save_results: 是否保存结果

        返回:
        完整预测结果DataFrame
        """
        try:
            self.output_progress("init", 0, "初始化完整房间占用预测系统")
            
            if not self.api_mode:
                print("完整房间占用预测系统")
                print("天气预报式房间占用预测")
                print("=" * 60)
                print(f"输入文件: {csv_file}")
                print(f"预测时长: {num_predictions*0.5:.1f}小时 ({num_predictions}个时间点)")
                print(f"使用模型: {models}")

            # 检查输入文件
            if not os.path.exists(csv_file):
                raise FileNotFoundError(f"输入文件不存在: {csv_file}")

            # 加载传感器模型
            self.load_sensor_model()

            # 读取历史数据
            self.output_progress("init", 16, "正在读取历史数据...")
            historical_data = pd.read_csv(csv_file)
            self.output_progress("init", 18, f"读取完成，共{len(historical_data)}条历史记录")

            # 第一阶段：预测传感器数据
            future_sensors_df = self.predict_future_sensors(historical_data, num_predictions)

            # 第二阶段：预测房间占用
            results_df = self.predict_room_occupancy(future_sensors_df, models)

            # 保存结果
            if save_results:
                self.output_progress("save", 92, "正在保存预测结果...")
                saved_files = self.save_results(results_df)

            # 准备结果数据
            output_file = saved_files[-1] if saved_files else None
            result_data = {
                "output_file": output_file,
                "num_predictions": len(results_df),
                "prediction_hours": len(results_df) * 0.5,
                "models_used": models,
                "summary": f"成功预测了未来{len(results_df)*0.5:.1f}小时的房间占用情况"
            }

            self.output_progress("complete", 100, "预测完成！")
            
            # 显示统计 (非API模式)
            if not self.api_mode:
                self.print_prediction_summary(results_df)

            # 输出最终结果
            self.output_result(True, result_data)
            return results_df

        except Exception as e:
            error_msg = str(e)
            self.output_result(False, error=error_msg)
            if not self.api_mode:
                import traceback
                traceback.print_exc()
            raise

def main():
    """主函数"""
    print("完整房间占用预测系统")
    print("天气预报式房间占用预测")

    # 获取输入参数
    print("\n请输入预测参数:")

    # 输入文件
    print("1. 使用默认数据文件 (room_occupancy.csv)")
    print("2. 输入自定义文件路径")
    choice = input("请选择 (1 或 2，默认为1): ").strip()

    if choice == '2':
        csv_file = input("请输入CSV文件路径: ").strip()
    else:
        csv_file = '../data/room_occupancy.csv'

    # 预测时长
    try:
        num_pred = input("预测时长选择 (1=24小时, 2=48小时, 3=自定义, 默认为2): ").strip()
        if num_pred == '1':
            num_predictions = 48  # 24小时
        elif num_pred == '3':
            custom_hours = float(input("请输入预测小时数: "))
            num_predictions = int(custom_hours * 2)  # 转换为30分钟间隔的点数
        else:
            num_predictions = 96  # 48小时
    except:
        num_predictions = 96

    # 模型选择
    print("\n选择使用的模型:")
    print("1. 使用所有模型 (RF + LSTM + XGBoost)")
    print("2. 仅使用随机森林 (RF)")
    print("3. 仅使用LSTM")
    print("4. 仅使用XGBoost")

    model_choice = input("请选择 (1-4，默认为1): ").strip()

    if model_choice == '2':
        models = ['rf']
    elif model_choice == '3':
        models = ['lstm']
    elif model_choice == '4':
        models = ['xgb']
    else:
        models = ['rf', 'lstm', 'xgb']

    # 运行完整预测
    try:
        system = CompleteRoomPredictionSystem()
        results = system.complete_prediction(
            csv_file=csv_file,
            num_predictions=num_predictions,
            models=models,
            save_results=True
        )

        print(f"\n预测完成！")
        print(f"成功预测了未来{num_predictions*0.5:.1f}小时的房间占用情况")

        # 显示部分结果
        print(f"\n预测结果示例 (前5行):")
        display_cols = ['Date', 'Time', 'S1_Temp', 'S5_CO2', 'ensemble_prediction']
        available_cols = [col for col in display_cols if col in results.columns]
        print(results[available_cols].head())

    except Exception as e:
        print(f"预测失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='完整房间占用预测系统')
    parser.add_argument('--input', required=True, help='输入CSV文件路径')
    parser.add_argument('--predictions', type=int, default=96, help='预测时间点数量')
    parser.add_argument('--models', nargs='+', default=['rf', 'lstm', 'xgb'], help='使用的模型')
    parser.add_argument('--api-mode', action='store_true', help='API模式，输出JSON格式的进度信息')
    
    args = parser.parse_args()
    
    if args.api_mode:
        # API模式
        try:
            system = CompleteRoomPredictionSystem(api_mode=True)
            results = system.complete_prediction(
                csv_file=args.input,
                num_predictions=args.predictions,
                models=args.models,
                save_results=True
            )
        except Exception as e:
            # 确保错误也通过正确的格式输出
            pass
    else:
        # 交互模式
        main()